package ai.langchain;

import dev.langchain4j.model.language.LanguageModel;
import dev.langchain4j.model.output.Response;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

public class OllamaDeepSeekModel implements LanguageModel {

    private static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";
    private static final String MODEL_NAME = "deepseek-r1:1.5b";
    private static final ObjectMapper MAPPER = new ObjectMapper();

    @Override
    public Response<String> generate(String prompt) {
        try (CloseableHttpClient client = HttpClients.createDefault()) {
            // 构造请求体
            String requestBody = String.format(
                "{\"model\": \"%s\", \"prompt\": \"%s\", \"stream\": false}",
                MODEL_NAME, prompt.replace("\"", "\\\"")
            );

            // 发送请求
            HttpPost post = new HttpPost(OLLAMA_API_URL);
            post.setHeader("Content-Type", "application/json");
            post.setEntity(new StringEntity(requestBody));

            try (CloseableHttpResponse response = client.execute(post)) {
                String jsonResponse = EntityUtils.toString(response.getEntity());
                JsonNode root = MAPPER.readTree(jsonResponse);
                
                // 解析响应
                if (root.has("response")) {
                    return Response.from(root.get("response").asText());
                } else {
                    throw new RuntimeException("API 响应异常: " + jsonResponse);
                }
            }
        } catch (Exception e) {
            throw new RuntimeException("模型调用失败", e);
        }
    }
}